import torch.nn as nn

from training.models.vit_mlp.mlp_mixer import mixer_b16_224_miil_in21k
from training.models.vit_mlp.unet import ModifiedUNet

__all__ = ['unet_mixer']


class UnetMixer(nn.Module):

    def __init__(self, num_classes, drop_rate, pretrained):
        super(UnetMixer, self).__init__()
        self.modified_unet = ModifiedUNet(3, 3)

        # self.mixer = gmlp_s16_224(num_classes=num_classes, drop_rate=drop_rate, pretrained=pretrained)
        self.mixer = mixer_b16_224_miil_in21k(num_classes=num_classes, drop_rate=drop_rate, pretrained=pretrained)


    def forward(self, x):
        spoof_cue = self.modified_unet(x)
        x = self.mixer(spoof_cue)
        return x, spoof_cue


def unet_mixer(num_classes=2, drop_rate=0., pretrained=False):
    model = UnetMixer(num_classes, drop_rate, pretrained)
    model.default_cfg = model.mixer.default_cfg
    return model
